#include "mem/mmap.h"
#include "mem/pmem.h"
#include "mem/vmem.h"
#include "proc/cpu.h"
#include "lib/print.h"
#include "lib/str.h"
#include "memlayout.h"

// 连续虚拟空间的复制(在uvm_copy_pgtbl中使用)
static void copy_range(pgtbl_t old, pgtbl_t new, uint64 begin, uint64 end)
{
    uint64 va, pa, page;
    int flags;
    pte_t* pte;

    for(va = begin; va < end; va += PGSIZE)
    {
        pte = vm_getpte(old, va, false);
        assert(pte != NULL, "uvm_copy_pgtbl: pte == NULL");
        assert((*pte) & PTE_V, "uvm_copy_pgtbl: pte not valid");
        
        pa = (uint64)PTE_TO_PA(*pte);
        flags = (int)PTE_FLAGS(*pte);

        page = (uint64)pmem_alloc(false);
        memmove((char*)page, (const char*)pa, PGSIZE);
        vm_mappages(new, va, page, PGSIZE, flags);
    }
}

// 两个 mmap_region 区域合并
// 保留一个 释放一个 不操作 next 指针
// 在uvm_munmap里使用
static void mmap_merge(mmap_region_t* mmap_1, mmap_region_t* mmap_2, bool keep_mmap_1)
{
    // 确保有效和紧临
    assert(mmap_1 != NULL && mmap_2 != NULL, "mmap_merge: NULL");
    assert(mmap_1->begin + mmap_1->npages * PGSIZE == mmap_2->begin, "mmap_merge: check fail");
    
    // merge
    if(keep_mmap_1) {
        mmap_1->npages += mmap_2->npages;
        mmap_region_free(mmap_2);
    } else {
        mmap_2->begin -= mmap_1->npages * PGSIZE;
        mmap_2->npages += mmap_1->npages;
        mmap_region_free(mmap_1);
    }
}

// 打印以 mmap 为首的 mmap 链
// for debug
void uvm_show_mmaplist(mmap_region_t* mmap)
{
    mmap_region_t* tmp = mmap;
    printf("\nmmap allocable area:\n");
    if(tmp == NULL)
        printf("NULL\n");
    while(tmp != NULL) {
        printf("allocable region: %p ~ %p\n", tmp->begin, tmp->begin + tmp->npages * PGSIZE);
        tmp = tmp->next;
    }
}

// 递归释放 页表占用的物理页 和 页表管理的物理页
// ps: 顶级页表level = 3, level = 0 说明是页表管理的物理页
static void destroy_pgtbl(pgtbl_t pgtbl, uint32 level) {
    // 确保页表级别有效
    assert(level <= 3, "destroy_pgtbl: Invalid level");

    // 遍历当前页表的每一项
    for (int i = 0; i < PGSIZE / sizeof(pte_t); i++) {
        pte_t page_entry = pgtbl[i]; // 获取当前页表项

        if (page_entry & PTE_V) { // 如果页表项有效
            if (level > 1) {
                if (!(PTE_CHECK(page_entry))) {
                    panic("destroy_pgtbl: Invalid page entry");
                }

                // 获取子页表的物理地址并递归释放
                pgtbl_t child_table = (pgtbl_t)PTE_TO_PA(page_entry);
                destroy_pgtbl(child_table, level - 1);
            } else {
                // 页表项指向物理页，直接释放
                pmem_free(PTE_TO_PA(page_entry), true);
            }

            // 清空页表项
            pgtbl[i] = 0;
        }
    }

    // 释放当前页表
    pmem_free((uint64)pgtbl, true);
}

// 页表销毁：trapframe 和 trampoline 单独处理
void uvm_destroy_pgtbl(pgtbl_t pgtbl)
{
    vm_unmappages(pgtbl, TRAPFRAME, PGSIZE, true);
    vm_unmappages(pgtbl, TRAMPOLINE, PGSIZE, false);
    destroy_pgtbl(pgtbl, 3);
}

// 拷贝页表 (拷贝并不包括trapframe 和 trampoline)
void uvm_copy_pgtbl(pgtbl_t old, pgtbl_t new, uint64 heap_top, uint32 ustack_pages, mmap_region_t *mmap) {
    // Step 1: 拷贝 USER_BASE 到 heap_top 范围的页表
    copy_range(old, new, USER_BASE, heap_top);

    // Step 2: 拷贝用户栈
    copy_range(old, new, TRAPFRAME - ustack_pages * PGSIZE, TRAPFRAME);

    // Step 3: 拷贝 mmap_region 的页表
    uint64 alloc_start = MMAP_BEGIN;
    for (mmap_region_t *region = mmap; region != NULL; region = region->next) {
        while (alloc_start < region->begin) {
            copy_range(old, new, alloc_start, alloc_start + PGSIZE);
            alloc_start += PGSIZE;
        }
    }
}

// 在用户页表和进程mmap链里 新增mmap区域 [begin, begin + npages * PGSIZE)
// 页面权限为perm
void uvm_mmap(uint64 begin, uint32 npages, int perm) {
    // 如果需要映射的页面数量为 0，直接返回
    if (npages == 0) return;

    // 确保起始地址对齐
    assert(begin % PGSIZE == 0, "uvm_mmap: start address not aligned");

    struct proc *p = myproc(); // 获取当前进程
    mmap_region_t *current = p->mmap; // 遍历的当前节点
    mmap_region_t *prev = NULL;      // 当前节点的前一个节点

    uint64 end = begin + npages * PGSIZE;

    // 遍历 mmap 链表，找到适合的空闲区域
    while (current) {
        uint64 region_end = current->begin + current->npages * PGSIZE;

        // 情况 1: 请求的内存区域完全位于当前区域中间，需拆分当前节点
        if (begin > current->begin && end < region_end) {
            mmap_region_t *new_region = mmap_region_alloc();
            if (!new_region) panic("uvm_mmap: failed to allocate mmap_region");

            // 分裂当前节点
            new_region->begin = end;
            new_region->npages = (region_end - end) / PGSIZE;
            current->npages = (begin - current->begin) / PGSIZE;

            // 插入新节点
            new_region->next = current->next;
            current->next = new_region;
            break;
        }

        // 情况 2: 请求的内存区域从当前节点的起始位置开始
        else if (begin == current->begin) {
            if (end < region_end) {
                current->begin += npages * PGSIZE;
                current->npages -= npages;
            } else if (end == region_end) {
                // 请求区域正好占满当前节点
                if (prev) {
                    prev->next = current->next;
                } else {
                    p->mmap = current->next;
                }
                mmap_region_free(current);
            }
            break;
        }

        // 情况 3: 请求的内存区域接在当前节点的末尾
        else if (end == region_end) {
            current->npages -= npages;
            break;
        }

        // 遍历下一个节点
        prev = current;
        current = current->next;
    }

    // 未找到合适的区域
    if (!current) panic("uvm_mmap: no suitable mmap_region found");

    // 完成页表映射，分配物理页
    uint64 va = begin;
    for (uint32 i = 0; i < npages; i++) {
        void *page = pmem_alloc(false);
        if (!page) panic("uvm_mmap: failed to allocate physical page");
        vm_mappages(p->pgtbl, va, (uint64)page, PGSIZE, perm);
        va += PGSIZE;
    }
}

// 在用户页表和进程mmap链里释放mmap区域 [begin, begin + npages * PGSIZE)
void uvm_munmap(uint64 begin, uint32 npages) {
    // 如果释放页面数量为 0，直接返回
    if (npages == 0) return;

    // 确保起始地址对齐
    assert(begin % PGSIZE == 0, "uvm_munmap: begin address not aligned");

    struct proc *proc = myproc();
    uint64 end = begin + npages * PGSIZE;

    // 创建一个新的 mmap_region
    mmap_region_t *new_region = mmap_region_alloc();
    if (!new_region) panic("uvm_munmap: failed to allocate new mmap_region");

    new_region->begin = begin;
    new_region->npages = npages;
    new_region->next = NULL;

    mmap_region_t *mp = proc->mmap;
    mmap_region_t *prev = NULL;

    // 遍历 mmap 链表并处理插入、合并的逻辑
    while (mp) {
        uint64 mp_end = mp->begin + mp->npages * PGSIZE;

        // 情况 1: 新区域位于链表第一个节点前，且不能合并
        if (end < mp->begin) {
            if (prev) {
                prev->next = new_region;
            } else {
                proc->mmap = new_region; // 更新头节点
            }
            new_region->next = mp;
            break;
        }

        // 情况 2: 新区域位于链表最后，且不能合并
        if (mp->next == NULL) {
            mp->next = new_region;
            new_region->next = NULL;
            break;
        }

        // 情况 3: 新区域接在当前节点后面
        if (mp_end == begin) {
            mmap_merge(mp, new_region, 1); // 合并当前节点与新区域
            if (end == mp->next->begin) { // 如果还能与下一节点合并
                mmap_merge(mp, mp->next, 1);
                mp->next = mp->next->next;
                mmap_region_free(mp->next); // 释放合并后的节点
            }
            break;
        }

        // 情况 4: 新区域接在当前节点前面
        if (end == mp->begin) {
            mmap_merge(new_region, mp, 0); // 合并新区域与当前节点
            if (prev) {
                prev->next = new_region;
            } else {
                proc->mmap = new_region; // 更新头节点
            }
            break;
        }

        // 情况 5: 新区域插入在两个节点之间，且不能合并
        if (begin > mp_end && end < mp->next->begin) {
            new_region->next = mp->next;
            mp->next = new_region;
            break;
        }

        // 移动到下一个节点
        prev = mp;
        mp = mp->next;
    }

    // 页表解除映射并释放物理页
    for (uint64 addr = begin; addr < end; addr += PGSIZE) {
        pte_t *pte = vm_getpte(proc->pgtbl, addr, 0);
        if (pte && (*pte & PTE_V)) {
            vm_unmappages(proc->pgtbl, addr, PGSIZE, 1); // 解除映射
        }
    }
}

// 用户堆空间增加
// 在这里无需修正 p->heap_top
//(注意栈顶最大值限制)
uint64 uvm_heap_grow(pgtbl_t pgtbl, uint64 heap_top, uint32 len, int flags)
{   
    // 计算新的堆顶
    uint64 new_heap_top = heap_top + len;

    // 校验是否超出用户地址空间上限
    // if (new_heap_top >= USER_STACK_TOP) {
    //     panic("uvm_heap_grow: heap exceeds stack limit.");
    // }

    // 页映射：批量分配
    uint64 va = heap_top; // 页对齐开始地址
    uint64 va_end = new_heap_top; // 页对齐结束地址

    for (; va < va_end; va += PGSIZE) {
        void *mem = pmem_alloc(false);
        if (!mem) {
            panic("uvm_heap_grow: memory allocation failed.");
        }
        memset(mem, 0, PGSIZE); // 清零内存
        vm_mappages(pgtbl, va, (uint64)mem, PGSIZE, PTE_U | flags);
    }

    // 返回新堆顶
    return new_heap_top;
}


// 用户堆空间减少
// 在这里无需修正 p->heap_top
uint64 uvm_heap_ungrow(pgtbl_t pgtbl, uint64 heap_top, uint32 len) {
    // 页对齐处理
    uint64 aligned_len = len;
    uint64 new_heap_top = heap_top - aligned_len;

    // 检查是否超出用户地址空间下限
    if (new_heap_top < USER_BASE) {
        panic("uvm_heap_ungrow: heap underflows user base.");
    }

    // 解映射页面
    uint64 va = new_heap_top; // 起始地址页对齐
    uint64 va_end = heap_top; // 结束地址页对齐

    if (va < va_end) {
        vm_unmappages(pgtbl, va, va_end - va, true); // 直接解映射范围
    }

    return new_heap_top;
}


// 用户态地址空间[src, src+len) 拷贝至 内核态地址空间[dst, dst+len)
// 注意: src dst 不一定是 page-aligned
void uvm_copyin(pgtbl_t pgtbl, uint64 dst, uint64 src, uint32 len) {
    while (len > 0) {
        // 计算 src 所在的页起始地址
        uint64 va0 = ALIGN_DOWN(src, PGSIZE);

        // 获取用户虚拟地址对应的物理地址
        pte_t *pte = vm_getpte(pgtbl, va0, false);
        if (!pte || (*pte & PTE_V) == 0) {
            panic("uvm_copyin: invalid user virtual address");
        }
        uint64 pa0 = PTE_TO_PA(*pte);

        // 计算当前页剩余可用字节数
        uint64 page_offset = src - va0;
        uint64 n = PGSIZE - page_offset;
        if (n > len) {
            n = len;
        }

        // 数据拷贝
        memmove((void *)dst, (void *)(pa0 + page_offset), n);

        // 更新指针和剩余长度
        src += n;
        dst += n;
        len -= n;
    }
}

// 内核态地址空间[src, src+len） 拷贝至 用户态地址空间[dst, dst+len)
void uvm_copyout(pgtbl_t pgtbl, uint64 dst, uint64 src, uint32 len) {
    while (len > 0) {
        // 计算 dst 所在的页起始地址
        uint64 va0 = ALIGN_DOWN(dst, PGSIZE);

        // 获取用户虚拟地址对应的物理地址
        pte_t *pte = vm_getpte(pgtbl, va0, false);
        if (!pte || (*pte & PTE_V) == 0) {
            panic("uvm_copyout: invalid kernel virtual address");
        }
        uint64 pa0 = PTE_TO_PA(*pte);

        // 计算当前页剩余可用字节数
        uint64 page_offset = dst - va0;
        uint64 n = PGSIZE - page_offset;
        if (n > len) {
            n = len;
        }

        // 数据拷贝
        memmove((void *)(pa0 + page_offset), (void *)src, n);

        // 更新指针和剩余长度
        src += n;
        dst += n;
        len -= n;
    }
}

// 用户态字符串拷贝到内核态
// 最多拷贝maxlen字节, 中途遇到'\0'则终止
// 注意: src dst 不一定是 page-aligned
void uvm_copyin_str(pgtbl_t pgtbl, uint64 dst, uint64 src, uint32 maxlen) {
    // uint32 copied = 0;
    while (maxlen > 0) {
        // 计算 src 所在的页起始地址
        uint64 va0 = ALIGN_DOWN(src, PGSIZE);

        // 获取用户虚拟地址对应的物理地址
        pte_t *pte = vm_getpte(pgtbl, va0, false);
        if (!pte || (*pte & PTE_V) == 0) {
            panic("uvm_copyin_str: invalid user virtual address");
        }
        uint64 pa0 = PTE_TO_PA(*pte);

        // 计算当前页剩余可用字节数
        uint64 page_offset = src - va0;
        uint64 n = PGSIZE - page_offset;
        if (n > maxlen) {
            n = maxlen;
        }

        // 字符拷贝，逐字节检查是否遇到 '\0'
        char *src_ptr = (char *)(pa0 + page_offset);
        char *dst_ptr = (char *)dst;
        while (n > 0) {
            *dst_ptr = *src_ptr;
            if (*src_ptr == '\0') {
                return; // 遇到 '\0' 提前返回
            }
            src_ptr++;
            dst_ptr++;
            src++;
            dst++;
            maxlen--;
            n--;
        }
    }

    // 如果未遇到 '\0' 且到达 maxlen，强制终止字符串
    *((char *)dst - 1) = '\0';
}
